package cn.doitedu.dwh

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator
import org.roaringbitmap.RoaringBitmap

object BitmapAggregator extends Aggregator[Array[Byte], Array[Byte], Array[Byte]] {
  // 初始化buff
  override def zero: Array[Byte] = {
    val buffBmp = RoaringBitmap.bitmapOf()
    BmpUtils.bmp2ByteArray(buffBmp)
  }

  // 局部聚合运算
  override def reduce(b: Array[Byte], a: Array[Byte]): Array[Byte] = {

    // 反序列化 buff数组  和输入的数组
    val buffBmp = BmpUtils.byteArray2Bmp(b)
    val inputBmp = BmpUtils.byteArray2Bmp(a)

    // 将输入的bitmap 或进  buff的bitmap
    buffBmp.or(inputBmp)

    // 将或好的结果序列化成字节数组，并返回
    BmpUtils.bmp2ByteArray(buffBmp)

  }

  // 全局聚合运算
  override def merge(b1: Array[Byte], b2: Array[Byte]): Array[Byte] = {
    reduce(b1,b2)
  }

  // 返回最终结果
  override def finish(reduction: Array[Byte]): Array[Byte] = reduction

  // 提供buff数据的encoder
  override def bufferEncoder: Encoder[Array[Byte]] = Encoders.BINARY

  // 提供返回结果的encoder
  override def outputEncoder: Encoder[Array[Byte]] = Encoders.BINARY
}
